{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac94ec5d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install rouge datasets==1.4.1 transformers==4.4.2 -i https://opentuna.cn/pypi/web/simple"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efb3646b",
   "metadata": {},
   "source": [
    "## preprocess data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8263ae8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "path = './demo_data/SUMMARY.meta'\n",
    "os.makedirs(path, exist_ok=True)\n",
    "print(\"The new directory is created!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "267f456c",
   "metadata": {},
   "outputs": [],
   "source": [
    "## data process \n",
    "import pandas as pd\n",
    "from datasets import load_dataset\n",
    "\n",
    "df_article_summary_full = pd.read_parquet('/home/ec2-user/SageMaker/CPT/finetune/generation/meta_description.parquet', engine='pyarrow')\n",
    "\n",
    "df_article_summary_full[['original_text','meta_descrption']].to_csv('./demo_data/SUMMARY.meta/total.csv',index=False)\n",
    "\n",
    "total_data = pd.read_csv('./demo_data/SUMMARY.meta/total.csv')\n",
    "x = total_data[-total_data['meta_descrption'].isnull()]\n",
    "x.columns = ['article','summarization']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "625758f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use csv file to test \n",
    "x[:1000].to_csv(os.path.join(path,'train.csv'),index=False,encoding='utf-8')\n",
    "x[1000:1200].to_csv(os.path.join(path,'test.csv'),index=False,encoding='utf-8')\n",
    "x[1200:1400].to_csv(os.path.join(path,'dev.csv'),index=False,encoding='utf-8')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62f2246e",
   "metadata": {},
   "source": [
    "## run train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a46b891",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "!python run_gen_v2.py --model_path 'fnlp/cpt-large' --dataset meta --data_dir demo_data --epoch '1' --batch_size '4' "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b90ac05",
   "metadata": {},
   "source": [
    "## 推理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbf754c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../')\n",
    "\n",
    "from transformers import BertTokenizer\n",
    "from modeling_cpt import CPTModel, CPTForConditionalGeneration\n",
    "\n",
    "model = CPTForConditionalGeneration.from_pretrained(\"output/meta/1\")\n",
    "tokenizer = BertTokenizer.from_pretrained('output/meta/1')\n",
    "dataset = load_dataset('csv', data_files='./demo_data/SUMMARY.meta/test.csv',split='train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6265368a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "def print_result(idx):\n",
    "    input_text = dataset[idx]['article']\n",
    "    print(\"input: \",input_text)\n",
    "\n",
    "    inputs = tokenizer(input_text, return_tensors=\"pt\",max_length=512)\n",
    "\n",
    "    #prompt_length = len(tokenizer.decode(inputs['input_ids'][0]))\n",
    "    #outputs = model.generate(inputs['input_ids'], max_length=64, do_sample=True, top_p=0.95, top_k=60)\n",
    "    outputs = model.generate(inputs['input_ids'], max_length=64, top_p=0.95)\n",
    "    generated = tokenizer.decode(outputs[0])\n",
    "\n",
    "    print(\"prediction result: \",generated)\n",
    "    \n",
    "    print ('label: ',dataset[idx]['summarization'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eb4125c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print_result(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d4c5ede",
   "metadata": {},
   "source": [
    "## 增强训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96dd4cea",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "!python run_gen_v2.py --model_path 'output/meta/1' --dataset meta --data_dir demo_data --epoch '1' --batch_size '4' "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f44bfa81",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_latest_p36",
   "language": "python",
   "name": "conda_pytorch_latest_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
